
import fitz  # PyMuPDF
import pytesseract
import cv2
import numpy as np
from PIL import Image
from typing import List, Tuple

DPI = 300        # 渲染分辨率，越高越准但越慢
TESS_LANG = 'chi_sim+eng'  # 中英双语

def render_page_to_cv2(page: fitz.Page) -> np.ndarray:
    """把 PDF 页面渲染成 OpenCV 格式 ndarray"""
    mat = fitz.Matrix(DPI/72, DPI/72)
    pix = page.get_pixmap(matrix=mat)
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

def ocr_with_table(image: np.ndarray) -> List[List[str]]:
    """Tesseract OCR 带表格结构；返回二维 list"""
    data = pytesseract.image_to_data(
        image,
        lang=TESS_LANG,
        output_type=pytesseract.Output.DATAFRAME,
        config='--psm 6 -c tessedit_create_tsv=1'
    )
    # 如果 Tesseract 没返回 table 结构，可以换用 camelot / paddleocr
    # 这里用最简：把 OCR 结果按行合并
    # 下面演示用空格对齐的简单做法，实际可再精细
    lines = []
    for _, g in data.groupby('line_num'):
        line = ' '.join(g.sort_values('left')['text'].dropna().astype(str))
        if line.strip():
            lines.append(line.strip())
    # 简单暴力：按空行拆成二维表
    table = [row.split() for row in lines if row.strip()]
    return table

def tables_with_titles(pdf_path: str) -> List[Tuple[str, List[List[str]]]]:
    doc = fitz.open(pdf_path)
    tables: List[Tuple[str, List[List[str]]]] = []

    for page in doc:
        img = render_page_to_cv2(page)
        # 1. OCR 整页文字，用于定位标题
        text = pytesseract.image_to_string(img, lang=TESS_LANG, config='--psm 6')
        lines = [l.strip() for l in text.splitlines() if l.strip()]

        # 2. OCR 表格区域
        table = ocr_with_table(img)

        if not table:
            continue

        # 3. 简单定位标题：表格第一行在整页 OCR 中的上一行
        title = "无标题"
        for i, line in enumerate(lines):
            if table[0] and table[0][0] in line:
                title = lines[i-1] if i > 0 else "无标题"
                break

        # 4. 合并逻辑（按标题+页码顺序即可）
        if tables and tables[-1][0] == title:
            tables[-1][1].extend(table[1:])  # 去掉表头
        else:
            tables.append((title, table))

    doc.close()
    return tables

# ------------------- 运行 -------------------
if __name__ == "__main__":
    for title, rows in tables_with_titles("aaa.pdf"):
        print("【标题】", title)
        for r in rows:
            print(r)
        print()
